"""
LongContextReorder 的主要目的是优化长文档在上下文中的排列顺序，以提高语言模型的处理效率和信息检索效果,特别适合需要处理大量文档或长文档的应用场景。

工作原理
问题背景：
当处理长文档时，语言模型(LLM)往往对上下文开头和结尾部分记忆更好
中间部分的信息容易被忽略或遗忘
传统顺序排列可能导致重要信息被埋没在文档中间

解决方案：
重新排列文档顺序，将最重要的信息放在开头和结尾
采用"反金字塔"结构：最重要的->次重要的->最不重要的->次重要的->重要的
"""
from pprint import pprint

from langchain_community.document_transformers import LongContextReorder
from langchain_community.vectorstores import Chroma

from models import get_ollama_embeddings_client

texts = [
    "Basquetball is a great sport.",
    "Fly me to the moon is one of my favourite songs.",
    "The Celtics are my favourite team.",
    "This is a document about the Boston Celtics",
    "I simply love going to the movies",
    "The Boston Celtics won the game by 20 points",
    "This is just a random text.",
    "Elden Ring is one of the best games in the last 15 years.",
    "L. Kornet is one of the best Celtics players.",
    "Larry Bird was an iconic NBA player.",
]

embeddings_client = get_ollama_embeddings_client()
retriever = Chroma.from_texts(texts, embedding=embeddings_client).as_retriever(search_kwargs={"k": 10})

query = "What can you tell me about the Celtics?"
docs = retriever.invoke(query)
print("重排前:")
pprint(docs)

reordering = LongContextReorder()
reordered_docs = reordering.transform_documents(docs)
print("重排后:")
# 将相似数据置于前后位置
pprint(reordered_docs)
